'''
  nuscenes数据集读取都是靠token标识,使用字典加载
  一个激光数据,6路摄像头,5路毫米波雷达
  1. 先获取token标识,
  2. 根据token获取数据
'''

from nuscenes.nuscenes import NuScenes
import numpy as np
import cv2
import os
from pyquaternion import Quaternion

version = "v1.0-mini"
dataroot = "/home/lin/code/bevdetv2.1/BEVDet/data/nuscenes"

nusc = NuScenes(version, dataroot, verbose=False)

# print("sample_number: ", len(nusc.sample), "\n")

samples = (nusc.sample)

timestamps = []
timestamps_src = []
for sample in samples:
    
    camera_token = sample['data']['CAM_FRONT']
    camera_sample_info = nusc.get("sample_data", camera_token)
    timestamp = camera_sample_info['timestamp']
    
    timestamps.append(int(str(timestamp)[:10]))
    timestamps_src.append(timestamp)
j = 0
k = 0
scene_id = 1  # 选场景

indexs = []

for i, time  in  enumerate(timestamps):
    if(i == 0):
        continue
    dist = time - timestamps[i-1]
    if(dist > 5):  # 统计2帧数据之间的时间戳，如果大于了2秒就，视频就不连续了
        k += 1
        print(k, i - j)   
        j = i
    if(k == scene_id):    # 统计有多少个场景， 并记录一个场景中所有index
        indexs.append(i)
    
print(indexs)
# cameras = ['CAM_FRONT', 'CAM_FRONT_LEFT', 'CAM_FRONT_RIGHT','CAM_BACK', 'CAM_BACK_LEFT', 'CAM_BACK_RIGHT']
cameras = ['CAM_FRONT']

for i in indexs:
    for cam in cameras:
        camera_token = samples[i]['data'][cam]
        camera_sample_info = nusc.get("sample_data", camera_token)
        camera_filename = os.path.join(dataroot, camera_sample_info['filename'])   
        print(camera_filename)
        img = cv2.imread(camera_filename)
        
        img = cv2.resize(img, [800, 450])
        cv2.imshow("src", img)
    
    if cv2.waitKey(100) == ord('q'):
        break